import json
import os
import cv2
import numpy as np
from collections import defaultdict


def create_my():
    data_dir = r'C:\Users\jiangzhengquan\Desktop\Switch_Dataset'
    img_path = os.path.join(data_dir, 'JPEGImages')
    with open(os.path.join(data_dir, 'train.txt'), 'w') as fp:
        for file in os.listdir(img_path):
            if file.endswith('.json'):
                name = os.path.splitext(file)[0]
                img = name + ".jpg"
                contents = json.load(open(os.path.join(img_path, file), 'r'))['shapes']
                path = img + " "
                for i, shape in enumerate(contents):
                    if i > 0:
                        path += " "
                    label = shape['label']
                    point = shape['points']
                    point = np.asarray(point, dtype=np.int32)
                    x1, y1, w, h = cv2.boundingRect(point)
                    x1, y1 = x1 - 30, y1 - 30
                    x2, y2 = x1 + w + 30, y1 + h + 30
                    path = path + ",".join(list(map(str, [x1, y1, x2, y2, int(label) - 1])))
                    # image = cv2.imread(os.path.join(img_path, name))
                    # cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 3)
                    # cv2.imshow('', image)
                    # cv2.waitKey(2)
                fp.write(path + '\n')


class COCO():
    def __init__(self, data_type='train'):
        super(COCO, self).__init__()
        self.data_dir = r'/home/jiangzhengquan/data/coco/annotations'
        self.annos = {}
        self.imgs = {}
        self.cats = {}
        self.data_type = data_type
        self.img2annos = defaultdict(list)
        self.img2cats = defaultdict(list)
        self.id_labels = {}
        self.names = []
        self.load_data()

        self.create_file()

    def load_data(self):
        train_file = os.path.join(self.data_dir, 'instances_{}2017.json'.format(self.data_type))
        train_data = json.load(open(train_file, 'r'))
        for anno in train_data['annotations']:
            self.annos[anno['id']] = anno
            self.img2annos[anno['image_id']].append(anno['id'])
            self.img2cats[anno['image_id']].append(anno['category_id'])
        for img in train_data['images']:
            self.imgs[img['id']] = img
        for cat in train_data['categories']:
            self.cats[cat['id']] = cat
            self.id_labels[cat['id']] = cat['name']
        with open('classes.txt', 'w') as fp:
            keys = self.id_labels.keys()
            keys = sorted(keys)
            for k in keys:
                print(k)
                name = self.id_labels[k]
                self.names.append(name)
                fp.write(name + '\n')
        self.label_id = dict((v, k) for k, v in enumerate(self.names))

    def create_file(self):
        with open(self.data_type + '.txt', 'w') as fp:
            for img_id, img_info in self.imgs.items():
                file_name = img_info['file_name']

                boxes = []
                labels = []
                for anno_id in self.img2annos[img_id]:
                    anno = self.annos[anno_id]
                    boxes.append(anno['bbox'])
                for cat_id in self.img2cats[img_id]:
                    cat = self.cats[cat_id]
                    n = cat['name']
                    labels.append(self.label_id[n])
                print(file_name)
                if len(boxes) == 0 or len(labels) == 0:
                    continue
                fp.write(file_name)
                for b, l in list(zip(boxes, labels)):
                    fp.write(' ')
                    b.append(l)
                    fp.write(','.join(list(map(str, b))))
                fp.write('\n')
                print(img_id)


if __name__ == '__main__':
    create_my()
